import os
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
from .train import GarbageClassifier

class GarbagePredictor:
    def __init__(self, model_path, device='cuda' if torch.cuda.is_available() else 'cpu'):
        self.device = device
        self.model = GarbageClassifier(num_classes=4)
        self.model.load_state_dict(torch.load(model_path, map_location=device))
        self.model.to(device)
        self.model.eval()
        
        self.classes = ['recyclable', 'harmful', 'wet', 'dry']
        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    
    def predict(self, image_path):
        # 加载并预处理图像
        image = Image.open(image_path).convert('RGB')
        image_tensor = self.transform(image).unsqueeze(0).to(self.device)
        
        # 进行预测
        with torch.no_grad():
            outputs = self.model(image_tensor)
            _, preds = torch.max(outputs, 1)
            pred_class = self.classes[preds.item()]
            
            # 计算各类别的概率
            probs = torch.nn.functional.softmax(outputs, dim=1)[0]
            class_probs = {cls: float(prob) for cls, prob in zip(self.classes, probs)}
        
        return {
            'predicted_class': pred_class,
            'probabilities': class_probs
        }

def main():
    # 示例用法
    model_path = 'best_model.pth'
    predictor = GarbagePredictor(model_path)
    
    # 测试图像预测
    test_image = 'path/to/test/image.jpg'
    if os.path.exists(test_image):
        result = predictor.predict(test_image)
        print(f"预测类别: {result['predicted_class']}")
        print("各类别概率:")
        for cls, prob in result['probabilities'].items():
            print(f"{cls}: {prob:.4f}")

if __name__ == '__main__':
    main()